using AnimCollections;
using AOT;
using Unity.Burst;
using UnityEngine.Animations;
using UnityEngine.Playables;

namespace DA.AnimGraph {
    [BurstCompile]
    public unsafe partial struct LayerBlendNode : INode {

        public struct LayerDesc {
            public bool isAdditive;
            public FloatLinkDesc weightDesc;
            public MaskLinkDesc maskDesc;
        }
        public struct SetupData : INodeSetupData {
            public LayerDesc[] layerDescs; // 不包括0端口
        }

        private struct LayerTempData {
            public bool isAdditive;
            public MaskLink mask;
        }

        //
        private NativeSimpleArray<FloatLink> linkWeights; // 不包括端口0，输入0的权重永远为1，所以这个数组的数量为输入数量减1
        // 临时数据，BuildPlayableGraph中可以删除，所以不放入buffer（划掉）
        // 但是！这样会影响 buffer 复制，所以还是放入buffer，可以不删除
        private NativeSimpleArray<LayerTempData> layerTempData;

        //
        private NodeHead head;
        private AnimationLayerMixerPlayable playable;


        [BurstCompile]
        [MonoPInvokeCallback(typeof(NodeLink.DisposeDelegate))]
        private static bool Dispose(in void* node) {
            var self = (LayerBlendNode*)node;

            self->linkWeights.Dispose();
            return true;
        }

        private void Setup(in SetupData data, ref GraphBuffer buffer, in NodeLink* link) {
            // 期望是 data.layerDescs 的元素数量刚好是 子节点数-1
            int childrenCountWithoutBase = data.layerDescs.Length;

            if (childrenCountWithoutBase > 0) {
                // 为了memcpy，还是要分配在 buffer 上
                //layerTempData = new NativeArray<LayerTempData>(data.layerDescs.Length, Allocator.Persistent);
                buffer.MallocNativeSimpleArray(out layerTempData, childrenCountWithoutBase);

                buffer.MallocNativeSimpleArray(out linkWeights, childrenCountWithoutBase);
                // 跳过了0端口
                for (int i = 0; i < childrenCountWithoutBase; i++) {
                    linkWeights[i] = new FloatLink(data.layerDescs[i].weightDesc);
                    ref var t = ref layerTempData.GetRefAt(i);
                    t.isAdditive = data.layerDescs[i].isAdditive;
                    t.mask = new MaskLink(data.layerDescs[i].maskDesc);
                }
            }
        }

        private static Playable BuildPlayableGraph(in void* node, in NodeLink* link, in PlayableGraph graph, in BuildContext buildContext) {
            var self = (LayerBlendNode*)node;

            self->head = new NodeHead(link);
            self->head.weights[0] = 1;
            self->head.weights[self->head.childrenCount] = 1;

            int childrenCount = self->head.childrenCount;
            self->playable = AnimationLayerMixerPlayable.Create(graph, childrenCount);


            self->playable.ConnectInput(0, self->head.children[0]->BuildPlayableGraph(graph, buildContext), 0);
            for (int i = 1; i < childrenCount; i++) {
                self->playable.ConnectInput(i, self->head.children[i]->BuildPlayableGraph(graph, buildContext), 0);
                self->playable.SetLayerAdditive((uint)i, self->layerTempData[i - 1].isAdditive);
                var mask = self->layerTempData[i - 1].mask.Get(buildContext);
                if (mask != null) {
                    self->playable.SetLayerMaskFromAvatarMask((uint)i, mask);
                }
            }

            // 可以不删除，以后还能copy
            //self->layerTempData.Dispose();

            return self->playable;
        }

        [BurstCompile]
        [MonoPInvokeCallback(typeof(NodeLink.EvaluateDelegate))]
        private static void Evaluate(in void* node, bool onEnter, in EvaluateContext evaluateContext, in GraphContext graphContext, ref UpstreamData upstream, ref ScopeContext scope) {
            var self = (LayerBlendNode*)node;

            {
                self->head.weights[0] = 1f;

                var childScope0 = new ScopeContext(scope, 1);
                NodeLink.Evaluate(self->head.children[0], evaluateContext, graphContext, ref upstream, ref childScope0);
                scope.Fusion(childScope0);
            }

            for (int i = 1; i < self->head.childrenCount; i++) {
                ref FloatLink alphaLink = ref self->linkWeights.GetRefAt(i - 1);
                float weight = alphaLink.Evaluate(onEnter, evaluateContext, graphContext, scope);
                self->head.weights[i] = weight;

                if (weight > 0) {
                    var childScope = new ScopeContext(scope, weight);
                    NodeLink.Evaluate(self->head.children[i], evaluateContext, graphContext, ref upstream, ref childScope);
                    scope.Fusion(childScope);
                }
            }
        }

        [BurstCompile]
        [MonoPInvokeCallback(typeof(NodeLink.PostEvaluateDelegate))]
        private static void PostEvaluate(in void* node, in EvaluateContext evaluateContext, in GraphContext graphContext, ref UpstreamData upstream, ref PostScopeContext postScope) {
            var self = (LayerBlendNode*)node;

            // todo:这里的规则需要调整
            // 可能需要跑所有非0权重的输入，在混合时
            float remainingWeight = 1;
            for (int i = self->head.childrenCount - 1; i >= 1; --i) {
                var weight = self->head.weights[i];
                var oldWeight = weight;
                remainingWeight = remainingWeight.Reduce(ref weight);
                var dtWeight = oldWeight - weight;
                if (dtWeight.IsNearZero()) continue;

                PostScopeContext childContex = new(ref postScope, dtWeight);
                NodeLink.PostEvaluate(self->head.children[i], evaluateContext, graphContext, ref upstream, ref childContex);
                postScope.Fusion(ref childContex);

                if (remainingWeight <= 0) break;
            }
            if (remainingWeight > 0) {
                PostScopeContext childContex = new(ref postScope, remainingWeight);
                NodeLink.PostEvaluate(self->head.children[0], evaluateContext, graphContext, ref upstream, ref childContex);
                postScope.Fusion(ref childContex);
            }
        }
    }
}
